{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Installation"
      ],
      "metadata": {
        "id": "qeaeXdqkFCZD"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's start by building a basic chatbot.For conversation history, we require a CloudSQL instance running in a Google\n",
        "Cloud project that the user or service account that executes your LangChain\n",
        "code is authenticated with. Look at the Google Cloud documentation for more\n",
        "information about [setting up and connecting to CloudSQL instances](https://cloud.google.com/sql/docs/mysql/connect-instance-cloud-shell).\n",
        "\n",
        "These libraries contain the core components of LangChain, as well as the ability to\n",
        "handle memory in a local MySQL database"
      ],
      "metadata": {
        "id": "pd3D45ERFGtH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Base components for using LangChain on Google Cloud\n",
        "! pip install -qU langchain-google-vertexai\n",
        "# LangGraph enables creation of graph-based chat agents\n",
        "! pip install -qU langgraph httpx\n",
        "# Grandalf is a tool to visualize graphs (like LangGraph)\n",
        "! pip install -qU grandalf # for visualizing the graph\n",
        "# Enables LangGraph to use SQLite for memory handling\n",
        "! pip install -qU langgraph-checkpoint-sqlite # for chatbot memory"
      ],
      "metadata": {
        "id": "wUqtoWT8wuyh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bcabe299-58f2-455a-da35-88b5662d903e"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m50.4/50.4 kB\u001b[0m \u001b[31m1.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m88.6/88.6 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m130.5/130.5 kB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.4/76.4 kB\u001b[0m \u001b[31m4.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.0/78.0 kB\u001b[0m \u001b[31m3.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m404.4/404.4 kB\u001b[0m \u001b[31m18.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m295.8/295.8 kB\u001b[0m \u001b[31m15.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m141.9/141.9 kB\u001b[0m \u001b[31m6.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.5/54.5 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m108.7/108.7 kB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m378.0/378.0 kB\u001b[0m \u001b[31m19.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.8/41.8 kB\u001b[0m \u001b[31m2.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# If you are using Google Colab, you need to restart your kernel after installation\n",
        "# You can skip this step in Vertex AI Workbench\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n"
      ],
      "metadata": {
        "id": "EWukDcrbSRaA"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# If you are running this notebook on Google Colab, run the cell below to authenticate your environment\n",
        "# If you are not using Google Colab, ensure that your Jupyter Notebook instance is connected to Google Cloud\n",
        "# See https://cloud.google.com/sdk/gcloud for more information\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ],
      "metadata": {
        "id": "hITWjJcqSa9_"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Set your GCP project and region\n",
        "\n",
        "import vertexai\n",
        "\n",
        "PROJECT_ID = \"my-project-id\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ],
      "metadata": {
        "id": "tt8qxqnDFgbc"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Import the class to interact with a Google Cloud Vertex AI Chatbot and set up a gemini-flash model\n",
        "\n",
        "from langchain_google_vertexai import ChatVertexAI\n",
        "\n",
        "model = ChatVertexAI(model=\"gemini-1.5-flash\")"
      ],
      "metadata": {
        "id": "yFFgUiojF-TI"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's start with a simple conversation. Conversations are always an exchange of `HumanMessage` and `AIMessage`. The `HumanMessage` is provided by the user, while the `AIMessage` is generated by the model. When you print the full `AIMessage`, you will see that it comes with a lot of metadata, such as safety and usage information."
      ],
      "metadata": {
        "id": "Pw6ZZFonUw47"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# First, we send a HumanMessage\n",
        "from langchain_core.messages import HumanMessage\n",
        "\n",
        "model.invoke([HumanMessage(content=\"Hello, my name is Max!\")])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D-D8aDybQeTH",
        "outputId": "06a8f1b6-2cd4-4b6f-fb71-8f5e925e0e53"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "AIMessage(content=\"Hello Max! 👋  It's nice to meet you. What can I do for you today? 😊 \\n\", additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [{'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}], 'usage_metadata': {'prompt_token_count': 7, 'candidates_token_count': 24, 'total_token_count': 31, 'cached_content_token_count': 0}, 'finish_reason': 'STOP'}, id='run-69df1672-0746-42f1-b5cc-3c222253017c-0', usage_metadata={'input_tokens': 7, 'output_tokens': 24, 'total_tokens': 31})"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Unfortunately, our chatbot does not have memory, and will not remember previous messages\n",
        "from langchain_core.messages import HumanMessage\n",
        "\n",
        "model.invoke([HumanMessage(content=\"What is my name?\")])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gjl7byrhTTua",
        "outputId": "b113b33f-2cb5-4a0a-cefc-1ab1ad24c11f"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "AIMessage(content=\"As an AI, I do not have access to personal information about you, including your name. \\n\\nTo know your name, you'll have to tell me! 😊 \\n\", additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [{'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}], 'usage_metadata': {'prompt_token_count': 5, 'candidates_token_count': 37, 'total_token_count': 42, 'cached_content_token_count': 0}, 'finish_reason': 'STOP'}, id='run-d05bfd55-d05c-4d14-9689-a451f27d7511-0', usage_metadata={'input_tokens': 5, 'output_tokens': 37, 'total_tokens': 42})"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# To work around this, we can manually send the conversation history to the bot\n",
        "from langchain_core.messages import HumanMessage, AIMessage\n",
        "\n",
        "model.invoke([\n",
        "    HumanMessage(content=\"Hello, my name is Max!\"),\n",
        "    AIMessage(content=\"Hello Max, it's a pleasure to meet you. How can I assist you today?\"),\n",
        "    HumanMessage(\"What is my name?\")\n",
        "    ])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5eTf49MHUUFB",
        "outputId": "27b09522-d76f-4f66-f059-cceb9d80c255"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "AIMessage(content='You told me your name is Max! 😊 \\n', additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [{'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}, {'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability_label': 'NEGLIGIBLE', 'blocked': False, 'severity': 'HARM_SEVERITY_NEGLIGIBLE'}], 'usage_metadata': {'prompt_token_count': 31, 'candidates_token_count': 11, 'total_token_count': 42, 'cached_content_token_count': 0}, 'finish_reason': 'STOP'}, id='run-9fb794ee-09e1-4046-84bb-05e0281262e4-0', usage_metadata={'input_tokens': 31, 'output_tokens': 11, 'total_tokens': 42})"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Sending the full message conversation on every invocation is cumbersome and requires a lot of coding work on your part. Luckily, Langchain have released Langgraph, which handles memory using the checkpoint class."
      ],
      "metadata": {
        "id": "SAcLRtktVO_B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Import necessary langgraph classes to define the chatbot and handle memory for you\n",
        "from langchain_core.prompts import ChatPromptTemplate\n",
        "from langchain_google_vertexai import ChatVertexAI\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "from langgraph.checkpoint.sqlite import SqliteSaver\n",
        "from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver\n",
        "from langgraph.graph import StateGraph\n",
        "from langgraph.graph.message import add_messages\n",
        "\n",
        "from typing import Annotated\n",
        "\n",
        "from typing_extensions import TypedDict\n",
        "\n",
        "# Create the agent\n",
        "\n",
        "class State(TypedDict):\n",
        "    # update add_messages to append messages not overwrite\n",
        "    messages: Annotated[list, add_messages]\n",
        "\n",
        "class StateMachine:\n",
        "\n",
        "    def __init__(self, memory):\n",
        "\n",
        "        model = ChatVertexAI(model=\"gemini-1.5-pro\")\n",
        "\n",
        "        # Define a system prompt that defines how the agent should interact with users\n",
        "        system_prompt = ChatPromptTemplate.from_messages(\n",
        "        [\n",
        "            (\n",
        "                \"system\",\n",
        "                \"\"\"You are an expert advisor in a bank. Your name is Terry.\n",
        "                You advise users on the three loan products of the bank:\n",
        "                1/ Short term loan, 6 month duration, 5% interest rate, up to USD 10000. Great for making quick payments.\n",
        "                2/ Mid-term loan, 24 month duration, 3% interest rate, up to USD 20000. Great for buying used cars.\n",
        "                3/ Long-term loan, 60 month duration, 2% interest rate, up to USD 50000. Great for buying a new car or a cheap house.\n",
        "                Only respond to user questions on these loans.\n",
        "                    \"\"\"\n",
        "            ),\n",
        "            (\"placeholder\", \"{messages}\"),\n",
        "            (\"human\", \"{input}\"),\n",
        "        ]\n",
        "        )\n",
        "\n",
        "        # Bind the system_prompt to the model\n",
        "        sally = system_prompt | model\n",
        "\n",
        "        def chatbot(state: State):\n",
        "            return {\"messages\": [sally.invoke(state[\"messages\"])]}\\\n",
        "\n",
        "        graph_builder = StateGraph(State)\n",
        "\n",
        "        graph_builder.add_node(\"chatbot\", chatbot)\n",
        "\n",
        "        # Set a finish point. This instructs the graph \"any time this node is run, you can exit.\"\n",
        "        graph_builder.set_finish_point(\"chatbot\")\n",
        "        graph_builder.set_entry_point(\"chatbot\")\n",
        "\n",
        "        # If you use memory, you need to provide a thread id to the chatbot\n",
        "        self.thread={\"configurable\": {\"thread_id\": \"2\"}}\n",
        "        self.chain=graph_builder.compile(checkpointer=memory)\n",
        "\n",
        "    def respond(self,USER_MESSAGE:str):\n",
        "        result=self.chain.invoke({\"messages\": (\"user\",USER_MESSAGE)},self.thread)\n",
        "\n",
        "        # Need to cut off the first 80 characters, which say \"AI Message\"\n",
        "        return result[\"messages\"][-1].pretty_repr()[80:]\n",
        "\n",
        "\n",
        "\n",
        "# Create memory for the bot\n",
        "with SqliteSaver.from_conn_string(\":memory:\") as memory:\n",
        "  # Initiate the memory\n",
        "  chatbot = StateMachine(memory)\n",
        "\n",
        "  # Draw the graph\n",
        "  chatbot.chain.get_graph().print_ascii()\n",
        "\n",
        "  print(\"Agent will keep going until you say 'STOP'\")\n",
        "  question = \"\"\n",
        "  while \"STOP\" not in question:\n",
        "      question = input(\"User: \")\n",
        "      answer = chatbot.respond(question)\n",
        "      print(\"Bot: \" + answer)\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0lANkoBMUroC",
        "outputId": "81336257-9a76-43d9-f505-2840ec944621"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "+-----------+  \n",
            "| __start__ |  \n",
            "+-----------+  \n",
            "      *        \n",
            "      *        \n",
            "      *        \n",
            " +---------+   \n",
            " | chatbot |   \n",
            " +---------+   \n",
            "      *        \n",
            "      *        \n",
            "      *        \n",
            " +---------+   \n",
            " | __end__ |   \n",
            " +---------+   \n",
            "Agent will keep going until you say 'STOP'\n",
            "User: Hey, I would like to order some pizza!\n",
            "Bot: \n",
            "\n",
            "I am sorry, I can only provide information on loan products offered by our bank.  Would you like to know more about our short-term, mid-term, or long-term loan options?\n",
            "User: STOP\n",
            "Bot: \n",
            "\n",
            "Okay, I understand. Please let me know if you have any questions about our loan products in the future.\n"
          ]
        }
      ]
    }
  ]
}